import torch
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
from torch.utils.data.dataset import Dataset
from glob import glob
import cv2
import numpy as np
from torchvision import transforms
from torch.utils.data import DataLoader
from Covariance_Pooling_Steganalytic_Network import Net
import logging
class MyDataset(Dataset):
  def __init__(self, cover_dir, stego_dir, transform=None):
    self.transform = transform
    
    self.cover_dir = cover_dir
    self.stego_dir = stego_dir

    self.cover_list = [x.split('/')[-1] for x in glob(self.cover_dir+'/*')]
    assert len(self.cover_list) != 0, "cover_dir is empty"
    
  def __len__(self):
    return len(self.cover_list)

  def __getitem__(self, idx):
    file_index = int(idx)

    cover_path=os.path.join(self.cover_dir,self.cover_list[file_index])
    stego_path=os.path.join(self.stego_dir,self.cover_list[file_index])
    
    cover_data = cv2.imread(cover_path, -1)
    stego_data = cv2.imread(stego_path, -1)
    # print(cover_data.shape, stego_data.shape)

    data = np.stack([cover_data, stego_data])
    label = np.array([0, 1], dtype='int32')

    sample = {'data': data, 'label': label}

    if self.transform:
      sample = self.transform(sample)

    return sample

class ToTensor():
  def __call__(self, sample):
    data, label = sample['data'], sample['label']

    # data = np.expand_dims(data, axis=1)
    data = data.astype(np.float32)
    # data = data / 255.0

    new_sample = {
      'data': torch.from_numpy(data),
      'label': torch.from_numpy(label).long(),
    }

    return new_sample
  
def main(checkpoint):
    eval_transform = transforms.Compose([
    ToTensor()
    ])  
    kwargs = {'num_workers': 1, 'pin_memory': True}
    BATCH_SIZE = 32
    # models = [
    #     ('A', '/home/zzc/hack/hack-2/Pretrained_En_models/A_ete/dataset/val/cover', '/home/zzc/hack/hack-2/Pretrained_En_models/A_ete/dataset/val/stego'),
    #     ('B', '/home/zzc/hack/hack-2/Pretrained_En_models/B_hps/dataset/val/cover', '/home/zzc/hack/hack-2/Pretrained_En_models/B_hps/dataset/val/stego'),
    #     ('C', '/home/zzc/hack/hack-2/Pretrained_En_models/C_stegnet/dataset/val/cover', '/home/zzc/hack/hack-2/Pretrained_En_models/C_stegnet/dataset/val/stego'),
    #     ('D', '/home/zzc/hack/hack-2/Pretrained_En_models/D_cci/dataset/val/cover', '/home/zzc/hack/hack-2/Pretrained_En_models/D_cci/dataset/val/stego'),
    #     ('E', '/home/zzc/hack/hack-2/Pretrained_En_models/E_unet/dataset/val/cover', '/home/zzc/hack/hack-2/Pretrained_En_models/E_unet/dataset/val/stego'),
    #     ('F', '/home/zzc/hack/hack-2/Pretrained_En_models/F_dca/dataset/val/cover', '/home/zzc/hack/hack-2/Pretrained_En_models/F_dca/dataset/val/stego'),
    #     ('G', '/home/zzc/hack/hack-2/Pretrained_En_models/G_ete_dcva/dataset/val/cover', '/home/zzc/hack/hack-2/Pretrained_En_models/G_ete_dcva/dataset/val/stego'),
    #     ('H', '/home/zzc/hack/hack-2/Pretrained_En_models/H_hinet/dataset/val/cover', '/home/zzc/hack/hack-2/Pretrained_En_models/H_hinet/dataset/val/stego'),
    #     ('I', '/home/zzc/hack/hack-2/Pretrained_En_models/I_resnet/dataset/val/cover', '/home/zzc/hack/hack-2/Pretrained_En_models/I_resnet/dataset/val/stego'),
    #     ('J', '/home/zzc/hack/hack-2/Pretrained_En_models/J_udh/dataset/val/cover', '/home/zzc/hack/hack-2/Pretrained_En_models/J_udh/dataset/val/stego')
    # ]

    models = [
        ('A', '/home/zzc/hack/hack-2/Pretrained_En_models/A_ete/dataset/val/secret', '/home/zzc/hack/hack-2/Pretrained_En_models/A_ete/dataset/val/secret_rev'),
        ('B', '/home/zzc/hack/hack-2/Pretrained_En_models/B_hps/dataset/val/secret', '/home/zzc/hack/hack-2/Pretrained_En_models/B_hps/dataset/val/secret_rev'),
        ('C', '/home/zzc/hack/hack-2/Pretrained_En_models/C_stegnet/dataset/val/secret', '/home/zzc/hack/hack-2/Pretrained_En_models/C_stegnet/dataset/val/secret_rev'),
        ('D', '/home/zzc/hack/hack-2/Pretrained_En_models/D_cci/dataset/val/secret', '/home/zzc/hack/hack-2/Pretrained_En_models/D_cci/dataset/val/secret_rev'),
        ('E', '/home/zzc/hack/hack-2/Pretrained_En_models/E_unet/dataset/val/secret', '/home/zzc/hack/hack-2/Pretrained_En_models/E_unet/dataset/val/secret_rev'),
        ('F', '/home/zzc/hack/hack-2/Pretrained_En_models/F_dca/dataset/val/secret', '/home/zzc/hack/hack-2/Pretrained_En_models/F_dca/dataset/val/secret_rev'),
        ('G', '/home/zzc/hack/hack-2/Pretrained_En_models/G_ete_dcva/dataset/val/secret', '/home/zzc/hack/hack-2/Pretrained_En_models/G_ete_dcva/dataset/val/secret_rev'),
        ('H', '/home/zzc/hack/hack-2/Pretrained_En_models/H_hinet/dataset/val/secret', '/home/zzc/hack/hack-2/Pretrained_En_models/H_hinet/dataset/val/secret_rev'),
        ('I', '/home/zzc/hack/hack-2/Pretrained_En_models/I_resnet/dataset/val/secret', '/home/zzc/hack/hack-2/Pretrained_En_models/I_resnet/dataset/val/secret_rev'),
        ('J', '/home/zzc/hack/hack-2/Pretrained_En_models/J_udh/dataset/val/secret', '/home/zzc/hack/hack-2/Pretrained_En_models/J_udh/dataset/val/secret_rev')
    ]
    
    model = Net().cuda()
  
    model.load_state_dict(checkpoint['original_state'])
    model.eval()
    for model_data in models:
        valid_dataset = MyDataset(model_data[1], model_data[2], eval_transform)
        valid_loader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False, **kwargs)

        correct = 0
        with torch.no_grad():
            for sample in valid_loader:
                data, label = sample['data'], sample['label']

                shape = list(data.size())
                data = data.reshape(shape[0] * shape[1], shape[4], *shape[2:4])
                label = label.reshape(-1)

                data, label = data.cuda(), label.cuda()

                output = model(data)
                pred = output.max(1, keepdim=True)[1]
                correct += pred.eq(label.view_as(pred)).sum().item()

        accuracy = correct / (len(valid_loader.dataset) * 2)
        print('{}Eval accuracy: {:.4f}'.format(model_data[0],accuracy))

if __name__ == '__main__':
  
  models_name = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J']
  checkpoint = None
  for name in models_name:
      path = '/home/zzc/hack/Cov-Pooling-Steganalytic-Network/secret_rev_save_models/e_' + name + '_d_covnet.pt'
      checkpoint = torch.load(path)
      print('________________' + name + '__________________')
      main(checkpoint)
      print('______________________________________________')